"""
算法的主要思路如下:

将负荷按需求从大到小排序,依次满足每个负荷的需求。
对于每个负荷,找到与其相连的线路和发电机。
按剩余容量从大到小排序与负荷相连的发电机。
从剩余容量最大的发电机开始供电,直到满足负荷需求或发电机达到出力上限。
计算线路流量,发电机流出的电量等于其出力,负荷流入的电量等于其需求。
这个算法尽可能平衡发电机和线路的负载,优先使用剩余容量大的发电机和线路,避免个别发电机或线路负载过高。算法的时间复杂度为O(n log n),其中n为负荷数量。

请注意,这个算法是一个简化的模型,实际电力系统的运行还需要考虑更多因素,如线路损耗、电压控制等。但这个算法可以作为一个基础,帮助理解电力系统的运行原理和优化思路。
"""
import numpy as np
from tqdm import tqdm
from utils.node import Node
import pickle
# 定义电力系统的通用模型
class PowerSystem:
    def __init__(self, generators, loads, lines_capacity):
        """
        初始化电力系统
        :param generators: dict, 发电机名称到其最大出力的映射
        :param loads: dict, 负荷名称到其需求的映射
        :param lines_capacity: dict, 线路的容量映射，形式为 {(节点1, 节点2): 容量}
        """
        self.generators = generators  # 发电机出力
        self.loads = loads  # 负荷需求
        self.lines_capacity = lines_capacity  # 线路容量
        self.flows = {line: 0 for line in lines_capacity}  # 初始化流量分配为0
        self.nodes = lines_capacity.copy()  # 剩余容量
        self.nodes.update(loads)
        self.nodes.update(generators)

    def find_min_load_path(self, source, target, visited:dict, max_load_rate=0.8):
        """
        使用DFS寻找从source到target的路径，优先选择负载较小的路径
        :param source: 源节点
        :param target: 目标节点
        :param visited: 已访问节点
        :return: 路径和路径的最小剩余容量
        """
        if source.get_id() in visited:
            return visited[source.get_id()]
        if source == target:
            return [[[target.get_id()],np.inf]]
        visited[source.get_id()] = [] # 防止环路
        paths = []
        # DEV 是否找到路线就跳出
        done = False
        # UNEND
        for node_id in self.nodes[source.get_id()].get_to_list():
            capacity = source.get_capacity()
            path = self.find_min_load_path(self.nodes[node_id], target, visited)
            if len(path)>0:  # 如果找到了路径
                for idx_, path_ in enumerate(path):
                    max_load = max(self.loads[path_[0][-1]]["load"]) # 负载节点的最大负载值
                    if  capacity >= max_load:
                        paths.append([[source.get_id()] + path_[0],min(capacity,path_[1])])
            # DEV 是否找到路线就跳出
            #             if source.get_load_rate(max_load)<max_load_rate:
            #                 done = True
            # if done:
            #     break
            # END
        visited[source.get_id()] = paths
        return paths

    def _find_load_path(self, source, target, visited:dict):
        """
        使用DFS寻找从source到target的路径，优先选择负载较小的路径
        :param source: 源节点
        :param target: 目标节点
        :param visited: 已访问节点
        :return: 路径和路径的最小剩余容量
        """
        if source.get_id() in visited:
            return visited[source.get_id()]
        if source == target:
            return [[target.get_id()]]
        visited[source.get_id()] = [] # 防止环路
        paths = []
        for node_id in self.nodes[source.get_id()].get_to_list():
            path = self._find_load_path(self.nodes[node_id], target, visited)
            if len(path)>0:  # 如果找到了路径
                for path_ in path:
                    paths.append([source.get_id()] + path_)
        visited[source.get_id()] = paths
        return paths

    def find_load_path(self, rdf_id):
        load = self.nodes[rdf_id]
        paths = {}
        for generator_id, generator in self.generators.items():
            if generator.get_capacity() > 0:  # 如果发电机还有剩余电力
                paths_ = self._find_load_path(generator, load, {}) # 所有的线路以及线路的剩余容量
                paths[generator_id] = paths_
                
        return paths

    def allocate_power(self):
        """
        分配电力，满足负荷需求并平衡线路负载
        """
        over_loads = []
        for rdf_id, load in tqdm(self.loads.items()): 
            # 对于每个负荷，找到一个从发电机到负荷的路径，并分配电力
            demand = max(load["load"])
            while demand > 0:
                best_path = None
                best_capacity = 0
                # 在所有发电机中寻找可行路径
                sort_generators = sorted(self.generators.values(), key = lambda x:x.get_load_rate())
                done = False
                for generator in sort_generators:
                    if generator.get_capacity() > 0:  # 如果发电机还有剩余电力
                        paths = self.find_min_load_path(generator, load, {}) # 所有的线路以及线路的剩余容量
                        # 遍历可行路径, 找出容量最大的路径
                        for path_ in paths:
                            if path_[1] > best_capacity:
                                best_path = path_[0]
                                best_capacity = path_[1]
                                done = True
                                break
                    if done:
                        break
                if best_path is None:
                    over_loads.append(rdf_id)
                    # print(f"{rdf_id} 无法满足负荷需求，系统容量不足。")
                    break

                # 分配电力，选择路径上容量最小的线路
                allocated_power = min(best_capacity, demand)

                # 更新路径上每一段线路的流量和剩余容量, 除了最终负载节点
                for idx_ in best_path[:-1]:
                    self.nodes[idx_].add_load(rdf_id,allocated_power)

                # 更新发电机的剩余电力和负荷的剩余需求
                demand -= allocated_power
                load["best_path"] = best_path
        with open('nodes_data.pickle', 'wb') as file:
            pickle.dump(self.nodes, file)
        print(over_loads)
        print("done")

    def print_flows(self):
        """
        打印电力流分配情况
        """
        print("电力流分配情况:")
        for line, flow in self.flows.items():
            print(f"{line}: {flow} MW")


# 测试通用电力系统
if __name__ == "__main__":

    from utils.cus_data_io import buid_test_data, load_data
    # generators, loads, lines = buid_test_data()
    generators, loads, lines = load_data()
    
    # 初始化电力系统
    power_system = PowerSystem(generators, loads, lines)
    # power_system.find_load_path("WINDING_151002856137513")

    # 分配电力
    power_system.allocate_power()

    # 输出电力流分配结果
    power_system.print_flows()